import requests
import json
import pickle
from pathlib import Path

from rich import print
from rich.console import Console
from rich.markdown import Markdown

# 初始化全局变量
DEEPSEEK_API_KEY = "sk-8e69e81224ab4567867d3351553d4ca8"  # 中医
# DEEPSEEK_API_KEY = "sk-a8369b74125a4c2ab9ee3a919e9e2da9"  # 处方
API_ENDPOINT = "https://api.deepseek.com/chat/completions"
MAX_HISTORY = 200  # 最大历史记录条数（防溢出）
SESSION_FILE_NAME = 'aicoder.pkl'
INIT_DIR ='aicoder'
# 初始化对话历史（系统消息可自定义）
# conversation_history = [
#     # {"role": "system", "content": "你是一位编程高手，尤其擅长python"}
# ]

def chat_with_deepseek(user_input: str, conversation):

    # 添加用户新输入到历史
    conversation.append({"role": "user", "content": user_input})

    # 创建请求头
    headers = {
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
        "Content-Type": "application/json"
    }

    # 创建请求体
    # payload = {
    #     "model": "deepseek-chat",
    #     "messages": conversation_history,
    #     "temperature": 0.0,  # 控制创造性（0-1）
    #     "max_tokens": 2000,  # 最大生成token数
    #     "stream": False      # 关闭流式传输（简化处理）
    # }

    payload = {
        "model": "deepseek-reasoner",
        "messages": conversation,
        "temperature": 0.4,  # 关键控制点（推荐范围0.1-0.3）
        "top_p": 0.5,  # 严格限制候选词范围
        "frequency_penalty": 0.4,  # 减少术语重复
        "presence_penalty": 0.2,  # 保持专业术语一致性
        "max_tokens": 10000  # 确保完整解释
        # "stream": True
    }

    try:
        # 发送API请求
        response = requests.post(
            API_ENDPOINT,
            headers=headers,
            data=json.dumps(payload),
            timeout=1000  # 设置超时时间（秒）
        )

        # 检查响应状态
        if response.status_code == 200:
            # 解析响应数据
            response_data = response.json()
            ai_reply = response_data['choices'][0]['message']['content']

            # 添加AI回复到历史（关键步骤！）
            conversation.append({
                "role": "assistant",
                "content": ai_reply
            })

            # 控制历史记录长度（防止超出token限制）
            if len(conversation) > MAX_HISTORY:
                # 保留最近的对话（移除最早的非系统消息）
                conversation = [conversation[0]] + conversation[-MAX_HISTORY+1:]

            return ai_reply

        else:
            error_msg = f"API错误: {response.status_code} - {response.text}"
            # conversation.append({
            #     "role": "system",
            #     "content": f"⚠️ 请求失败: {error_msg}"
            # })
            return f"请求失败，请稍后重试。错误详情: {error_msg}"

    except Exception as e:
        error_msg = f"网络异常: {str(e)}"
        # conversation.append({
        #     "role": "system",
        #     "content": f"⚠️ 网络错误: {str(e)}"
        # })
        return f"网络连接异常: {str(e)}"


def save_session(filename, conversation):
    with open(filename, 'wb') as f:
        pickle.dump(conversation, f)


def load_init_files(folder, conversation):

    folder_path = Path(folder)
    for file in folder_path.iterdir():
        if file.is_file():  # 确保是文件而不是文件夹
            with open(file, 'r', encoding='utf-8') as f:
                conversation.append({
                    "role": "user",
                    "content": f.read()
                })


def load_session(filename, init_dir, system_first):
    conversation = []
    try:
        with open(filename, 'rb') as f:
            conversation.extend(pickle.load(f))
    except FileNotFoundError:
        # conversation_history = [{"role": "system", "content": "会话文件找不到，不能保存当前对话"}]
        print("从初始化目录中读取需求...")
        if system_first is not None:
            conversation.append({"role": "system", "content": system_first})
        load_init_files(init_dir, conversation)

    # 可选：打印当前历史记录（调试用）
    print("\n[当前对话历史]")
    for msg in conversation:
        print(msg['role'])
        print(Markdown(f"{msg['content']}"))
        print()

    return conversation


def start(sessionb_file_name, init_dir, system_first=None):
    print("DeepSeek 对话开始（输入'exit'退出）")
    conversation_history = load_session(sessionb_file_name, init_dir, system_first)

    # if system_first is not None:
    #     conversation_history.append({"role": "system", "content": system_first})
    # load_init_files(init_dir, conversation_history)
    # load_session(sessionb_file_name, init_dir, conversation_history)

    while True:
        user_input = input("\nuser: ")

        if user_input.lower() in ['exit', 'quit', 'e', 'q']:
            print("对话结束")
            # save_session(sessionb_file_name, conversation_history)
            break

        if user_input.lower() in ['pop']:
            print("去掉最后一次对话，并退出")
            conversation_history.pop()
            conversation_history.pop()
            save_session(sessionb_file_name, conversation_history)
            break

        # 获取AI回复
        response = chat_with_deepseek(user_input, conversation_history)
        print(f"\nassistant:")
        print(Markdown(f"{response}"))
        save_session(sessionb_file_name, conversation_history)


if __name__ == "__main__":
    start(SESSION_FILE_NAME)
